# coding: utf-8
import argparse
import json
import os
import re
import time

import jsonlines

from LLMBase import LLMBase
from logger import Logger


def preprocess(content):
    """
    remove img element in markdown: ![.*?](.*?)
    remove table element, its in html format
    replace multiple empty lines with single one
    """
    content = re.sub(r"!\[.*?\]\(.*?\)", "", content)
    # content = re.sub(r"<table>.*?</table>", "", content, flags=re.DOTALL)
    content = re.sub(r"\n{2,}", "\n", content)
    return content


def save_qa_jsonl(questions, answers, raw_data, save_path):
    assert len(questions) == len(answers), "Questions and Answers are not in pairs"
    qa_pairs = []
    for q, a in zip(questions, answers):
        qa_pairs.append(
            {
                "Q": q,
                "A": a,
                "raw_data_id": raw_data["data_id"],
                "book_name": raw_data["bookname"],
                "chapter": raw_data["chapter"],
                "book_type": raw_data["book_type"],
            }
        )

    with jsonlines.open(save_path, "a") as writer:
        for qa in qa_pairs:
            writer.write(qa)


def save_qa(questions, answers):
    # assert len(questions) == len(answers), "Questions and Answers are not in pairs"
    qa_pairs = []
    for q, a in zip(questions, answers):
        qa_pairs.append({"Q": q, "A": a})

    with open("qa_pairs.json", "w", encoding="utf-8") as f:
        json.dump(qa_pairs, f, ensure_ascii=False, indent=2)


# with open("chapter.txt", "r", encoding="utf-8") as f:
#     book_content = preprocess(f.read())


def extract_questions(text: str):
    return [q.strip() for q in text.split("\n") if q.strip()]


def get_qa(llm, prompt):
    messages = [
        {"role": "user", "content": prompt},
    ]
    try:
        res = llm.chat(messages, sys_prompt="你是一名土木工程专业结构抗震课程的助教")
    except Exception as e:
        llm.logger.log(f"ERROR GET response from LLM, error: {e}")
        return [], []

    content = llm.record_response(res)
    if "none" in content.lower():
        return [], []

    questions = extract_questions(content)
    llm.logger.log(f"questions:\n" + "\n".join(questions))

    messages.append({"role": "assistant", "content": content})

    prefix = "现假设你是一名教授大学建筑结构抗震课程的资深教师，请根据材料给出这些问题的解答，要求每个问题的解答务必详实，对涉及到的概念展开阐述。注意，回答不要索引所给材料中的图、表等，回答应该能够在脱离该材料的独立语境中含义明确。涉及公式的内容应给出具体公式，而不是引用材料中的编号。尽量能分点回答，条理清晰。"
    answers = []
    for i in range(len(questions)):
        prompt = (
            prefix
            + f"请回答第{i+1}个问题:{questions[i]}。仅返回解答即可，请勿在回复末尾附加你的与解答无关的补充或解释性话语"
            if i == 0
            else f"请回答第{i+1}个问题:{questions[i]}。仅返回解答即可，请勿在回复末尾附加你的与解答无关的补充或解释性话语"
        )
        messages.append({"role": "user", "content": prompt})

        res = llm.chat(messages, sys_prompt="你是一名教授大学建筑结构抗震课程的资深教师")
        llm.logger.log(f"answer {i+1} finished, total tokens of this chat: {res.usage.total_tokens}")
        content = llm.record_response(res)
        answers.append(content)
        messages.append({"role": "assistant", "content": content})
    return questions, answers


def first_ask(llm, raw_data, question_num=3):
    prompt = f"""假设你是结构抗震课程的助教，请根据下面给出的材料中（某抗震教材中截取的章节内容）提出{question_num}个专业性问题，以考察对知识的思考和理解，问题的难易程度不限。注意问题应在没有该材料的情况下语义明确，即不要出现‘根据材料‘，’根据图x.xx’等话语。示例：Q1: 结构地震反应分析中常用的方法有哪些？Q2: 以建筑结构抗震设计方法为基础，请说明第一阶段设计和第二阶段设计的具体流程和目的，并指出哪些情况下需要进行第二阶段设计？如何确保达到全面满足抗震设防目标的要求？如果无法满足前述要求或提出有效问题，则返回None；若能提出，请按如下格式返回：1.<question1>\n2.<question2>...\n\n材料："""

    text = raw_data["content"]
    prompt += text
    questions, answers = get_qa(llm, prompt)
    return questions, answers


def repeat_ask(llm, raw_data, existing_questions):
    existing_questions = "\n".join(existing_questions)
    # print(existing_questions)
    prompt = f"假设你是结构抗震课程的助教，请根据下面给出的材料中（某抗震教材中截取的章节内容）提出3个专业性问题，以考察对知识的思考和理解，问题的难易程度不限。问题应在没有该材料的情况下语义明确，即不要出现‘根据材料‘，’根据图x.xx’等话语。现已有如下一些问题：\n{existing_questions}，\n提出的问题需避免和以上这些问题类似。如果无法再提出新的问题，则返回None；若能提出，请按如下格式返回：1.<question1>\n2.<question2>...\n\n材料："

    prompt += raw_data["content"]
    questions, answers = get_qa(llm, prompt)
    return questions, answers


def create_qa(llm, raw_data, save_path):
    all_q, all_a = [], []
    qs, aws = first_ask(llm, raw_data)
    if not qs:
        return
    save_qa_jsonl(qs, aws, raw_data, save_path)
    all_q.extend(qs)
    all_a.extend(aws)

    while True:
        qs, aws = repeat_ask(llm, raw_data, all_q)
        if not qs:
            print("No more questions can be asked, have generated", len(all_q), "questions")
            break
        save_qa_jsonl(qs, aws, raw_data, save_path)
        all_q.extend(qs)
        all_a.extend(aws)
        if len(all_q) >= 6:
            # print("Enough questions have been asked")
            break

    return all_q, all_a
    # save_qa(all_q, all_a)
    # print("Done")


def process_single_item(llm, raw_data_path, save_path, start_id, end_id):
    finished = []
    if os.path.exists(result_path):
        with jsonlines.open(result_path, "r") as reader:
            finished = [item["raw_data_id"] for item in reader.iter()]

    all_q, all_a = [], []
    with jsonlines.open(raw_data_path, "r") as reader:
        for item in reader:
            if item["data_id"] in finished:
                continue
            if start_id <= int(item["data_id"]) < end_id:
                print(">>> Processing item: ", item["data_id"])
                start = time.time()
                qs, aws = create_qa(llm, item, save_path)
                end = time.time()
                print("cost time: ", end - start)
                all_q.extend(qs)
                all_a.extend(aws)
    # return all_q, all_a


## config

model_name = "ChatGPT"
segment_path = "text_blocks_v2/books/segments.jsonl"
result_path = f"text_blocks_v2/books/qa_pairs_{model_name}.jsonl"
parser = argparse.ArgumentParser()
parser.add_argument("--start_index", type=int, default=0)
parser.add_argument("--total", type=int, default=100)
args = parser.parse_args()
# start running
logger = Logger(f"chat_scheme_{model_name}.txt")
llm = LLMBase(model_name, logger)
process_single_item(
    llm, segment_path, result_path, start_id=args.start_index, end_id=args.total + args.start_index
)
# save_qa(questions, answers)
